
#include "../iface/APRInterface.h"
#include "../iface/debug.h"
#include <apr_atomic.h>
#include <apr_file_io.h>
#include "TCPServer.h"
#include <new>

#define CONNECTION_TIMEOUT  5000
#define DATA_TIMEOUT        20000
#define POLL_TIMEOUT_USEC   5000

TCPServer::TCPServer(int inoofthreads, int iportno, const char *ihostname, apr_pool_t *upperpool)
{
   apr_status_t status;
   status = apr_pool_create(&mp, upperpool);
   noofthreads = inoofthreads;
   listenstop = 0;
   sessionstop = 0;
   portno = iportno;
   listenstop = 1;
   if (ihostname)
   {
      hostname = (char *)apr_palloc(mp, sizeof(char)*(strlen(ihostname)+1));
      strcpy(hostname, ihostname);
   } else hostname = NULL;
   pipearray = (ClientPipes *)apr_palloc(mp, sizeof(ClientPipes)*noofthreads);
   for (int i = 0; i < noofthreads; i++)
   {
      pipearray[i].recv_to_handler = NULL;
      pipearray[i].handler_to_send = NULL;
      void *mem_for_new_waiting_lock = apr_palloc(mp, sizeof(APRMutex));
      pipearray[i].waiting = new(mem_for_new_waiting_lock) APRMutex(mp);
      void *mem_for_new_sessionstart_lock = apr_palloc(mp, sizeof(APRMutex));
      pipearray[i].sessionstart = new(mem_for_new_sessionstart_lock) APRMutex(mp);
      pipearray[i].stage = SERVERSTAGE_INACTIVE;
      pipearray[i].localmp = NULL;
   }
   session_threads = (apr_thread_t **)apr_palloc(mp, sizeof(apr_thread_t *)*noofthreads);
   session_threadattrs = (apr_threadattr_t **)apr_palloc(mp, sizeof(apr_threadattr_t *)*noofthreads);
}

TCPServer::~TCPServer()
{
   JoinThreads();
   // don't do free on pipearray, session_threads
   apr_pool_destroy(mp);
}

int TCPServer::StartListeningThread()
{
   apr_status_t status;
   status = apr_threadattr_create(&listening_attr, mp);
   CheckAPRError(status);
   status = apr_thread_create(&listening_thread, listening_attr, (apr_thread_start_t)((functor)(&(TCPServer::ServerListen))), (void *)this, mp);
   CheckAPRError(status);
   return 0;
}

int TCPServer::StartSessionThreads()
{
   void ***tempvoid;
   tempvoid = (void ***)apr_palloc(mp, sizeof(void **)*noofthreads);
   apr_status_t status;
   for (int i = 0; i < noofthreads; i++)
   {
      tempvoid[i] = (void **)apr_palloc(mp, sizeof(void *)*2);
      tempvoid[i][0] = (void *)apr_palloc(mp, sizeof(void *));
      tempvoid[i][1] = (void *)apr_palloc(mp, sizeof(int));
      *((void **)tempvoid[i][0]) = (void *)this;
      *((int *)tempvoid[i][1]) = i;
      status = apr_threadattr_create(&(session_threadattrs[i]), mp);
      CheckAPRError(status);
      // Create stack protectors to prevent stack smashing attacks
      //status = apr_threadattr_guardsize_set(session_threadattrs[i], 4096);
      //CheckAPRError(status);
      status = apr_thread_create(&(session_threads[i]), session_threadattrs[i], (apr_thread_start_t)((functor)(GetClientHandlerFunction())), (void *)tempvoid[i], mp);
      CheckAPRError(status);
   }
   return 0;
}

void TCPServer::SocketAccept(apr_pollset_t *pollset, apr_socket_t *lsock, apr_pool_t *mp)
{
   // Must find inactive socket, or this does nothing (which is good!)
   apr_status_t status;
   for (int j = 0; j < noofthreads; j++)
   {
      if (apr_atomic_read32(&(pipearray[j].stage)) == SERVERSTAGE_INACTIVE)
      {
         if (pipearray[j].localmp != NULL)
         {
            pipearray[j].recv_to_handler->~APRBucketBrigade();
            pipearray[j].handler_to_send->~APRBucketBrigade();
            pipearray[j].recv_to_handler = NULL;
            pipearray[j].handler_to_send = NULL;
            apr_pool_clear(pipearray[j].localmp);
            apr_pool_destroy(pipearray[j].localmp);
         }
         status = apr_pool_create(&(pipearray[j].localmp), mp);
         CheckAPRError(status);
         void *mem_for_recv_to_handler = (APRBucketBrigade *)apr_palloc(mp, sizeof(APRBucketBrigade));
         void *mem_for_handler_to_send = (APRBucketBrigade *)apr_palloc(mp, sizeof(APRBucketBrigade));
         pipearray[j].recv_to_handler = new(mem_for_recv_to_handler) APRBucketBrigade(mp);
         pipearray[j].handler_to_send = new(mem_for_handler_to_send) APRBucketBrigade(mp);
         status = apr_socket_accept(&(pipearray[j].clisock), lsock, mp);
         CheckAPRError(status);
         status = apr_socket_addr_get(&(pipearray[j].cliaddr), APR_REMOTE, pipearray[j].clisock);
         CheckAPRError(status);
         apr_pollfd_t pfd = { pipearray[j].localmp, APR_POLL_SOCKET, APR_POLLIN, 0, { NULL }, NULL };
         pfd.desc.s = pipearray[j].clisock;
         status = apr_socket_opt_set(pipearray[j].clisock, APR_SO_NONBLOCK, 1);
         CheckAPRError(status);
         status = apr_socket_timeout_set(pipearray[j].clisock, DATA_TIMEOUT);
         CheckAPRError(status);
         status = apr_pollset_add(pollset, &pfd);
         CheckAPRError(status);
         apr_atomic_set32(&(pipearray[j].stage), SERVERSTAGE_RECEIVING);
         pipearray[j].sessionstart->signal();
         return;
      }
   }
}

void *TCPServer::ServerListen(apr_thread_t *tthread, void *vthis)
{
   apr_status_t status;
   TCPServer *kthis = (TCPServer *)vthis;
   apr_pool_t *upperpool, *mp;
   upperpool = apr_thread_pool_get(tthread);
   apr_pool_create(&mp, upperpool);

   apr_socket_t *lsock = kthis->CreateListeningSocket(mp);
   if (lsock == NULL)
   {
      kthis->listenstop = 1;
      return NULL;
   }
   apr_pollset_t *pollset;

   status = apr_pollset_create(&pollset, kthis->noofthreads, mp, 0);
   CheckAPRError(status);
   apr_pollfd_t mainpollfd = { mp, APR_POLL_SOCKET, APR_POLLIN, 0, { NULL }, NULL };
   const apr_pollfd_t *workingpollfd;
   mainpollfd.desc.s = lsock;
   status = apr_pollset_add(pollset, &mainpollfd);
   CheckAPRError(status);
   char *data = (char *)malloc(sizeof(char)*MAX_BLOCK_SIZE);

   while (!(kthis->listenstop))
   {
      apr_int32_t fdnum;
      // Using same timeout for polling as for connections!
      status = apr_pollset_poll(pollset, POLL_TIMEOUT_USEC, &fdnum, &workingpollfd);
      if (!(APR_STATUS_IS_EINTR(status) || APR_STATUS_IS_TIMEUP(status)))
         CheckAPRError(status);
      for (int i = 0; i < fdnum; i++)
      {
         int matched = 0;
         if (workingpollfd[i].desc.s == lsock)
         {
            kthis->SocketAccept(pollset, lsock, mp);
            matched = 1;
         } else
         {
            for (int j = 0; j < kthis->noofthreads; j++)
            {
               if (apr_atomic_read32(&(kthis->pipearray[j].stage)) > SERVERSTAGE_INACTIVE)
               {
                  if (kthis->pipearray[j].clisock == workingpollfd[i].desc.s)
                  {
                     if (kthis->pipearray[j].handler_to_send->IsClosed())
                     {
                        // Handle the premature close case
                        apr_pollfd_t pfd = { kthis->pipearray[j].localmp, APR_POLL_SOCKET, APR_POLLIN, 0, { NULL }, NULL };
                        pfd.desc.s = kthis->pipearray[j].clisock;
                        status = apr_pollset_remove(pollset, &pfd);
                        apr_atomic_set32(&(kthis->pipearray[j].stage), SERVERSTAGE_INACTIVE);
                        CheckAPRError(status);
                     }
                     if (workingpollfd[i].rtnevents & APR_POLLIN)
                     {
                        if (kthis->pipearray[j].recv_to_handler->PackedSpaceRemaining())
                        {
                           apr_size_t dataquant = MAX_BLOCK_SIZE;
                           status = apr_socket_recv(kthis->pipearray[j].clisock, data, &dataquant);
                           if (APR_STATUS_IS_EOF(status))
                           {
                              kthis->pipearray[j].recv_to_handler->Close();
                           } else
                           {
                              CheckAPRError(status);
                              if (status != 0)
                              {
                                 kthis->pipearray[j].recv_to_handler->Close();
                              }                                 
                           }
                           if (dataquant > 0)
                           {
                              /*int byteswritten = */kthis->pipearray[j].recv_to_handler->PackedWrite(data, dataquant);
                              kthis->pipearray[j].waiting->signal();
                           }
                        }
                     }
                     if (workingpollfd[i].rtnevents & APR_POLLOUT)
                     {
                        apr_size_t dataquant = kthis->pipearray[j].handler_to_send->Read(data, MAX_BLOCK_SIZE);
                        status = apr_socket_send(kthis->pipearray[j].clisock, data, &dataquant);
                        CheckAPRError(status);
                        kthis->pipearray[j].handler_to_send->Pop(data, dataquant);
                        apr_pollfd_t pfd = { kthis->pipearray[j].localmp, APR_POLL_SOCKET, APR_POLLIN|APR_POLLOUT, 0, { NULL }, NULL };
                        pfd.desc.s = kthis->pipearray[j].clisock;
                        status = apr_pollset_remove(pollset, &pfd);
                        CheckAPRError(status);
                        pfd.reqevents &= ~APR_POLLOUT;
                        status = apr_pollset_add(pollset, &pfd);
                        CheckAPRError(status);
                        kthis->pipearray[j].waiting->signal();
                     }
                     matched = 1;
                  }
               }
            }
         }
         if (matched == 0)
            errprint("Unmatched/unknown file descriptor in poll.");
      }
      for (int j = 0; j < kthis->noofthreads; j++)
      {
         if (apr_atomic_read32(&(kthis->pipearray[j].stage)) > SERVERSTAGE_INACTIVE)
         {
            if (kthis->pipearray[j].recv_to_handler->HasData())
               kthis->pipearray[j].waiting->signal();
            if (kthis->pipearray[j].handler_to_send->HasData())
            {
               apr_pollfd_t pfd = { kthis->pipearray[j].localmp, APR_POLL_SOCKET, APR_POLLIN, 0, { NULL }, NULL };
               pfd.desc.s = kthis->pipearray[j].clisock;
               status = apr_pollset_remove(pollset, &pfd);
               CheckAPRError(status);
               pfd.reqevents |= APR_POLLOUT;
               status = apr_pollset_add(pollset, &pfd);
               CheckAPRError(status);
            } else
            {
               if ((apr_atomic_read32(&(kthis->pipearray[j].stage)) >= SERVERSTAGE_SENDING) ||
                   (kthis->pipearray[j].handler_to_send->IsClosed()))
               {
                  if (apr_atomic_read32(&(kthis->pipearray[j].stage)) > SERVERSTAGE_SENDING)
                     errprint("Reset from unknown state (misunderstood data on connection)");
                  apr_pollfd_t pfd = { kthis->pipearray[j].localmp, APR_POLL_SOCKET, APR_POLLIN|APR_POLLOUT, 0, { NULL }, NULL };
                  pfd.desc.s = kthis->pipearray[j].clisock;
                  status = apr_pollset_remove(pollset, &pfd);
                  CheckAPRError(status);
                  status = apr_socket_close(kthis->pipearray[j].clisock);
                  CheckAPRError(status);
                  apr_atomic_set32(&(kthis->pipearray[j].stage), SERVERSTAGE_INACTIVE);
                  kthis->pipearray[j].waiting->signal();
               }
            }
         }
      }
   }
   free(data);
   kthis->listenstop = 1;
   return NULL;
}

void *TCPServer::ClientHandler(apr_thread_t *tthread, void *args)
{
   apr_status_t status;
   TCPServer *kthis = *((TCPServer **)(((void **)(args))[0]));
   int i = *((int *)(((void **)(args))[1]));
   apr_pool_t *upperpool, *mp;
   upperpool = apr_thread_pool_get(tthread);
   status = apr_pool_create(&mp, upperpool);
   CheckAPRError(status);

   while (!(kthis->sessionstop))
   {
      kthis->pipearray[i].sessionstart->wait();
      if (kthis->sessionstop) break;
      apr_pool_t *sp;
      apr_pool_create(&sp, mp);
      int has_read = 0, has_written = 0, has_finalised = 0;
      while ((apr_atomic_read32(&(kthis->pipearray[i].stage)) > SERVERSTAGE_INACTIVE) &&
             (apr_atomic_read32(&(kthis->pipearray[i].stage)) < SERVERSTAGE_SENDING))
      {
         kthis->pipearray[i].waiting->wait();
         if ((kthis->pipearray[i].recv_to_handler->IsClosed()) && (has_finalised == 0))
         {
            has_finalised = 1;
            apr_atomic_inc32(&(kthis->pipearray[i].stage));
         }
         char *buffer = (char *)malloc(sizeof(char)*MAX_BLOCK_SIZE);
         while ((kthis->pipearray[i].handler_to_send->PackedSpaceRemaining() >= MAX_BLOCK_SIZE) && 
                (kthis->pipearray[i].recv_to_handler->HasData()))
         {
            apr_size_t dataquant = kthis->pipearray[i].recv_to_handler->Pop(buffer, MAX_BLOCK_SIZE);
            if ((dataquant > 0) && (has_read == 0))
            {
               has_read = 1;
               apr_atomic_inc32(&(kthis->pipearray[i].stage));
            }
            apr_size_t dataquant2 = kthis->pipearray[i].handler_to_send->PackedWrite(buffer, dataquant);
            if ((dataquant2 > 0) && (has_written == 0))
            {
               has_written = 1;
               apr_atomic_inc32(&(kthis->pipearray[i].stage));
            }
         }
         free(buffer);
      }
      kthis->pipearray[i].waiting->unlock();
      apr_pool_clear(sp);
      apr_pool_destroy(sp);
   }
   kthis->pipearray[i].sessionstart->unlock();
   apr_pool_clear(mp);
   apr_pool_destroy(mp);
   ERR_remove_state(0);
   apr_thread_exit(tthread, APR_SUCCESS);
   return NULL;
}

int TCPServer::JoinThreads()
{
   sessionstop = 1;
   apr_status_t status;
   for (int i = 0; i < noofthreads; i++)
   {
      if (pipearray[i].sessionstart)
         pipearray[i].sessionstart->signal();
      apr_thread_join(&status, session_threads[i]);
      CheckAPRError(status);
   }
/*
   for (int i = 0; i < kthis->noofthreads; i++)
      pipearray[i].modify_mutex->signal();
*/

   listenstop = 1;
   apr_thread_join(&status, listening_thread);
   CheckAPRError(status);
   for (int i = 0; i < noofthreads; i++)
   {
      if (pipearray[i].localmp != NULL)
      {
         pipearray[i].recv_to_handler->~APRBucketBrigade();
         pipearray[i].handler_to_send->~APRBucketBrigade();
         pipearray[i].recv_to_handler = NULL;
         pipearray[i].handler_to_send = NULL;
         apr_pool_clear(pipearray[i].localmp);
         apr_pool_destroy(pipearray[i].localmp);
      }
   }
   return 0;
}

apr_socket_t *TCPServer::CreateListeningSocket(apr_pool_t *mp)
{
   apr_status_t status;
   apr_socket_t *tsocket;
   apr_sockaddr_t *addrstruct;

   // NULL here gives the 0.0.0.0 that is binding to all network interfaces
   status = apr_sockaddr_info_get(&addrstruct, hostname, APR_INET, portno, 0, mp);
   CheckAPRError(status);
   status = apr_socket_create(&tsocket, addrstruct->family, SOCK_STREAM, APR_PROTO_TCP, mp);
   CheckAPRError(status);
   status = apr_socket_opt_set(tsocket, APR_SO_NONBLOCK, 1);
   CheckAPRError(status);
   status = apr_socket_timeout_set(tsocket, CONNECTION_TIMEOUT);
   CheckAPRError(status);
   status = apr_socket_opt_set(tsocket, APR_SO_REUSEADDR, 1);
   CheckAPRError(status);
   status = apr_socket_bind(tsocket, addrstruct);
   CheckAPRError(status);
   if (status != APR_SUCCESS)
   {
      errprint("bind to '%s:%d' failed", hostname, portno);
      return NULL;
   }
   // makes the number of backlog connections the most the OS will support
   status = apr_socket_listen(tsocket, SOMAXCONN);
   CheckAPRError(status);

   listenstop = 0;
   return tsocket;
}

const char *TCPServer::StaticStringGetClientAddress(int threadid)
{
   apr_status_t status;
   char *ipaddress;
   status = apr_sockaddr_ip_get(&ipaddress, pipearray[threadid].cliaddr);
   CheckAPRError(status);
   return ipaddress;
}
